from common_import import *


def process_products():
    # 样例数据
    sales_data = tool.get_np("2x.csv")
    cost_data = tool.get_np("3x.csv")

    # 重点关注产品列表
    focused_products = mapping.focus_products

    # 日期范围
    start_date = datetime.date(2020, 7, 1)
    end_date = datetime.date(2023, 6, 30)
    date_range = np.arange(
        start_date, end_date + datetime.timedelta(days=1), dtype="datetime64[D]"
    )

    for product_code in focused_products:
        # 筛选出该产品的销售数据
        product_sales = sales_data[sales_data["code"] == product_code]

        # 筛选出该产品的成本价数据
        product_costs = cost_data[cost_data["code"] == product_code]

        # 创建一个字典用于存储每天的销售量总和和最大销售价格
        sales_dict = {}
        for row in product_sales:
            date = row["date"]
            quantity = row["quantity"]
            price = row["price"]
            if date in sales_dict:
                sales_dict[date] = (
                    sales_dict[date][0] + quantity,  # 累加当天的销售量
                    max(sales_dict[date][1], price),  # 取当天的最大价格
                )
            else:
                sales_dict[date] = (quantity, price)

        # 创建一个字典用于快速查找成本价数据
        cost_dict = {row["date"]: row["costprice"] for row in product_costs}

        # 初始化结构化数组，包括costprice字段
        structured_array = np.zeros(
            len(date_range),
            dtype=[
                ("date", "datetime64[D]"),
                ("quantity", "f4"),
                ("sales_price", "f4"),
                ("cost_price", "f4"),
            ],
        )

        # 填充结构化数组
        for i, date in enumerate(date_range):
            str_date = str(date)
            if str_date in sales_dict:
                quantity, price = sales_dict[str_date]
            else:
                quantity, price = 0, 0

            if str_date in cost_dict:
                costprice = cost_dict[str_date]
            else:
                costprice = 0

            structured_array[i] = (date, quantity, price, costprice)
        tool.get_csv(structured_array, f"problem3/{product_code}.csv")


def calculate_total_daily_profit():
    total_profit_dict = {}
    product_ids = mapping.focus_products
    for product_id in product_ids:
        # 提取商品数据
        data = tool.get_np(f"product_forcast/{product_id}.csv")

        # 计算每个商品的盈利
        profit = data["quantity"] * (data["sales_price"] - data["cost_price"])

        # 累加到每天的总盈利
        for i in range(len(data)):
            date = data["date"][i]
            if date in total_profit_dict:
                total_profit_dict[date] += profit[i]
            else:
                total_profit_dict[date] = profit[i]

    # 将总盈利转换为结构化数组
    result = np.zeros(
        len(total_profit_dict),
        dtype=[("date", "datetime64[D]"), ("total_profit", "f4")],
    )

    for i, (date, total_profit) in enumerate(sorted(total_profit_dict.items())):
        result[i] = (date, total_profit)

    return result[1:]


if __name__ == "__main__":
    pass
    result = calculate_total_daily_profit()
    print(result)
